import pickle
import numpy as np

vocab, embeddings = [], []

embedpath = r"D:\00_code\007_papars\chinese_vector\sgns.merge.word\sgns.merge.word"
outpath = r'D:\00_code\007_papars\NER-joint\joint\Recurrent_Interaction_Network_update\Recurrent_Interaction_Network_EMNLP2020-master\tcm_data'


word_vec_file = open(embedpath, 'r', encoding='utf8')

charlines = word_vec_file.readlines()
wordnum = None
embdim = None
worddic = {}

for i, charli in enumerate(charlines):
    if i == 0:
        wordnum = charli.strip().split()[0]
        embdim = int(charli.strip().split()[1])
        print(embdim)
    else:
        charli = charli.strip()
        if charli != "":
            try:
                word = charli.split()[0]
                wordemb = list(map(float, charli.split()[1:]))
                #pdb.set_trace()
                assert len(wordemb) == embdim, "the vector of char embedding is not match !"
                if worddic.get(word) is None:
                    worddic[word] = wordemb
                    vocab.append(word)
                    embeddings.append(wordemb)
                else:                           #重复的字符就只取第一个

                    print("repeat word : ", word)
            except:
                print(charli)

vocab_npa = np.array(vocab)
embs_npa = np.array(embeddings)

#insert '<pad>' and '<unk>' tokens at start of vocab_npa.
vocab_npa = np.insert(vocab_npa, 0, '<PAD>')
vocab_npa = np.insert(vocab_npa, 1, '<UNK>')
print(vocab_npa[:10])

pad_emb_npa = np.zeros((1,embs_npa.shape[1]))   #embedding for '<pad>' token.
unk_emb_npa = np.mean(embs_npa,axis=0,keepdims=True)    #embedding for '<unk>' token.

#insert embeddings for pad and unk tokens at top of embs_npa.
embs_npa = np.vstack((pad_emb_npa,unk_emb_npa,embs_npa))
print(embs_npa.shape)

with open(outpath+'/vocab.pkl', 'wb') as f:
    pickle.dump(vocab_npa, f)

# with open(outpath+'vocab_npa.npy','wb') as f:
#     np.save(f,vocab_npa)

with open(outpath+'/embedding.npy','wb') as f:
    np.save(f,embs_npa)